import numpy as np
import nibabel as nib


def npz_to_nii(npz_file_path, image_nii_path, mask_nii_path):
    """
    将 .npz 文件中的 image 和 mask 数据分别转换为 .nii 文件。

    参数:
    - npz_file_path: 输入 .npz 文件路径。
    - image_nii_path: 输出 image 的 .nii 文件路径。
    - mask_nii_path: 输出 mask 的 .nii 文件路径。

    返回:
    - None，直接生成 .nii 文件。
    """
    # 加载 .npz 文件
    npz_file = np.load(npz_file_path)

    # # 确认包含的键
    # print("文件包含的键:", npz_file.files)
    # if 'image' not in npz_file or 'mask' not in npz_file:
    #     raise ValueError("文件中缺少 'image' 或 'mask' 键，请检查文件内容")

    # 提取 image 和 mask 数据
    image = npz_file['image']
    mask = npz_file['label']
    print(f"加载 image 数据形状: {image.shape}, 数据类型: {image.dtype}")
    print(f"加载 mask 数据形状: {mask.shape}, 数据类型: {mask.dtype}")

    # 检查维度并调整为 3D 或 4D
    if image.ndim == 2:
        image = image[np.newaxis, ...]  # 添加深度维度
    if mask.ndim == 2:
        mask = mask[np.newaxis, ...]  # 添加深度维度

    if image.ndim > 4 or mask.ndim > 4:
        raise ValueError("数据维度超过 4，请检查输入数据")

    # 创建仿射矩阵 (4x4)
    affine = np.eye(4)  # 单位矩阵

    # 创建并保存 image 的 NIfTI 文件
    image_nifti = nib.Nifti1Image(image, affine)
    nib.save(image_nifti, image_nii_path)
    print(f"Image 保存成功: {image_nii_path}")

    # 创建并保存 mask 的 NIfTI 文件
    mask_nifti = nib.Nifti1Image(mask, affine)
    nib.save(mask_nifti, mask_nii_path)
    print(f"Mask 保存成功: {mask_nii_path}")


# 示例使用
npz_file_path = "/home/sust/mf/MISSFormer-main/datas/Spine/train/case0008_slice006.npz"  # 输入 .npz 文件路径
image_nii_path = "/home/sust/mf/MISSFormer-main/datas/Spine/nii/image/case0008_slice006.nii"  # 输出 image 的 .nii 文件路径
mask_nii_path = "/home/sust/mf/MISSFormer-main/datas/Spine/nii/mask/case0008_slice006.nii"  # 输出 mask 的 .nii 文件路径

npz_to_nii(npz_file_path, image_nii_path, mask_nii_path)
